"""工具服务 - 参数修改和降雨假拟的业务逻辑"""

import json
import os
from typing import Dict, Any, Optional
from pathlib import Path

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv

from mcp.forecast_client import ForecastAPIClient
from mcp.forecast_analyzer import ForecastAnalyzer
from .parameter_ranges import ParameterRanges


# 加载环境变量
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(env_path, override=True)


class Config:
    """服务配置"""
    # vLLM
    VLLM_BASE_URL = os.getenv("VLLM_BASE_URL", "http://localhost:8000/v1")
    VLLM_MODEL_NAME = os.getenv("VLLM_MODEL_NAME", "Qwen/Qwen2.5-3B-Instruct")
    VLLM_API_KEY = os.getenv("VLLM_API_KEY", "EMPTY")
    
    @classmethod
    def print_config(cls):
        print(f"🔧 [tools_service] vLLM 配置:")
        print(f"   VLLM_BASE_URL: {cls.VLLM_BASE_URL}")
        print(f"   VLLM_MODEL_NAME: {cls.VLLM_MODEL_NAME}")


class APIResponse:
    """统一 API 响应构建器"""
    
    @staticmethod
    def success_response(
        message: str,
        intent_type: str,
        understood: bool = True,
        **kwargs
    ) -> Dict[str, Any]:
        """创建成功响应"""
        return {
            "success": True,
            "message": message,
            "understood": understood,
            "intent_type": intent_type,
            **kwargs
        }
    
    @staticmethod
    def error_response(
        message: str,
        intent_type: str,
        errors: Optional[list] = None,
        understood: bool = False,
        **kwargs
    ) -> Dict[str, Any]:
        """创建错误响应"""
        return {
            "success": False,
            "message": message,
            "understood": understood,
            "intent_type": intent_type,
            "validation_errors": errors,
            "current_parameters": {},
            **kwargs
        }


class ToolsService:
    """工具服务类 - 处理参数修改和降雨假拟"""
    
    def __init__(self):
        """初始化服务"""
        self.llm = ChatOpenAI(
            base_url=Config.VLLM_BASE_URL,
            model_name=Config.VLLM_MODEL_NAME,
            api_key=Config.VLLM_API_KEY,
            temperature=0,
            max_tokens=500,
        )
    
    async def modify_parameters(self, question: str, session_id: str) -> Dict[str, Any]:
        """
        参数修改工具
        
        功能：
        - 解析用户的参数修改请求
        - 验证参数范围（基于 DHF 模型）
        - 生成参数修改 API 请求
        
        Args:
            question: 用户问题，如 "把SA0改成50" 或 "把K增加0.3"
            session_id: 会话ID
            
        Returns:
            符合 API 文档格式的字典
        """
        # Step 1: 获取当前参数值
        try:
            client = ForecastAPIClient()
            data = await client.fetch_forecast_data(session_id)
            analyzer = ForecastAnalyzer(data)
            current_params = analyzer.get_model_parameters()
            
            current_initial_state = current_params.get("initial_state", {})
            current_model_params = current_params.get("parameters", {})
            
            # 检查是否获取到参数
            if not current_initial_state and not current_model_params:
                return APIResponse.error_response(
                    message="无法获取会话参数，请检查会话ID是否正确",
                    intent_type="parameter_modification",
                    errors=["无法获取会话参数"],
                    update_success=False,
                    update_message="无法获取会话参数"
                )
        except Exception as e:
            return APIResponse.error_response(
                message=f"获取参数失败: {str(e)}",
                intent_type="parameter_modification",
                errors=[str(e)],
                update_success=False,
                update_message=f"获取参数失败: {str(e)}"
            )
        
        # Step 2: LLM 直接计算最终参数值
        extraction_prompt = ChatPromptTemplate.from_messages([
            ("system", """你是参数计算专家。根据用户请求和当前参数值，直接计算出最终的参数值。

当前参数值：
{current_params}

用户可能说：
- "把SA0改成50" 或 "把sa0改成50" → SA0最终值是50
- "把K增加0.3" 或 "把k增加0.3" → 如果当前K是0.5，最终K是0.8
- "把B减少0.1" 或 "把b减少0.1" → 如果当前B是2.0，最终B是1.9
- "a加上0.1" → 如果a参数对应A，如果当前A是1.0，最终A是1.1

请直接返回最终的参数值，JSON格式，只返回JSON。
重要：参数名必须使用大写字母（如 SA0, K, B, A 等）。

示例格式：
{{
  "SA0": 50,
  "K": 0.8
}}

只返回用户要修改的参数，不要返回所有参数。"""),
            ("human", "{question}")
        ])
        
        chain = extraction_prompt | self.llm
        response = await chain.ainvoke({
            "question": question,
            "current_params": json.dumps({
                "initial_state": current_initial_state,
                "parameters": current_model_params
            }, ensure_ascii=False, indent=2)
        })
        
        # 解析 LLM 返回的最终参数值
        try:
            content = response.content.strip()
            # 去除可能的 markdown 标记
            if "```" in content:
                content = content.split("```")[1]
                if content.startswith("json"):
                    content = content[4:]
            
            final_params = json.loads(content.strip())
        except json.JSONDecodeError:
            return APIResponse.error_response(
                message="无法解析参数修改请求",
                intent_type="parameter_modification",
                errors=["LLM返回格式错误"],
                explanation=f"LLM 返回格式错误: {response.content[:100]}",
                current_parameters={"initial_state": current_initial_state, "parameters": current_model_params},
                update_success=False,
                update_message="无法解析参数修改请求"
            )
        
        # Step 3: 验证范围并分类
        updates = {
            "initial_state": {},
            "parameters": {}
        }
        parameter_info = {}
        validation_errors = []
        
        for param_name, new_value in final_params.items():
            # 标准化参数名（支持小写输入）
            normalized_param_name = ParameterRanges.normalize_param_name(param_name)
            
            # 获取参数范围
            param_range = ParameterRanges.get_range(normalized_param_name)
            if not param_range:
                validation_errors.append(f"⚠️  {param_name} 不是已知的 DHF 模型参数")
                continue
            
            # 确定参数类型（使用标准化的参数名）
            if normalized_param_name in ParameterRanges.INITIAL_STATE_RANGES:
                param_category = "initial_state"
                current_value = current_initial_state.get(normalized_param_name, 0)
            else:
                param_category = "parameters"
                current_value = current_model_params.get(normalized_param_name, 0)
            
            # 验证范围
            min_val = param_range["min"]
            max_val = param_range["max"]
            
            # 记录参数信息（使用标准化的参数名）
            parameter_info[normalized_param_name] = {
                "name": param_range["name"],
                "current": current_value,
                "new": new_value,
                "min": min_val,
                "max": max_val,
                "unit": param_range["unit"],
                "in_range": min_val <= new_value <= max_val
            }
            
            # 范围检查
            if not parameter_info[normalized_param_name]["in_range"]:
                validation_errors.append(
                    f"❌ {param_range['name']} ({normalized_param_name}) 的值 {new_value} "
                    f"超过了范围 [{min_val}, {max_val}]"
                )
            else:
                # 添加到修改列表（使用标准化的参数名）
                updates[param_category][normalized_param_name] = new_value
        
        # Step 4: 组织返回数据（按照 API 文档格式）
        valid_count = sum(1 for info in parameter_info.values() if info["in_range"])
        
        # 构建 parameter_changes（使用路径格式：initial_state.SA0）
        parameter_changes = {}
        for param_name, info in parameter_info.items():
            if info["in_range"]:
                if param_name in ParameterRanges.INITIAL_STATE_RANGES:
                    parameter_changes[f"initial_state.{param_name}"] = info["new"]
                else:
                    parameter_changes[f"parameters.{param_name}"] = info["new"]
        
        # 构建 parameter_ranges
        parameter_ranges = {}
        for param_name, info in parameter_info.items():
            if param_name in ParameterRanges.INITIAL_STATE_RANGES:
                parameter_ranges[f"initial_state.{param_name}"] = [info["min"], info["max"]]
            else:
                parameter_ranges[f"parameters.{param_name}"] = [info["min"], info["max"]]
        
        # 生成消息
        if valid_count > 0:
            # 成功修改
            change_desc = []
            for param_name, info in parameter_info.items():
                if info["in_range"]:
                    change_desc.append(f"{info['name']}从 {info['current']} 修改为 {info['new']}")
            message = f"好的，我已将 {', '.join(change_desc)}"
            
            explanation_parts = []
            for param_name, info in parameter_info.items():
                if info["in_range"]:
                    explanation_parts.append(f"将{info['name']} {param_name} 修改为 {info['new']}")
            explanation = f"根据您的要求，{', '.join(explanation_parts)}"
            
            return APIResponse.success_response(
                message=message,
                intent_type="parameter_modification",
                parameter_changes=parameter_changes,
                explanation=explanation,
                current_parameters={
                    "initial_state": {**current_initial_state, **updates["initial_state"]},
                    "parameters": {**current_model_params, **updates["parameters"]}
                },
                validation_errors=validation_errors if validation_errors else None,
                update_success=True,
                update_message="参数验证通过",
                update_data={
                    "session_id": session_id,
                    "updated_parameters": list(parameter_changes.keys())
                },
                parameter_ranges=parameter_ranges
            )
        else:
            # 所有参数都超出范围
            return APIResponse.success_response(
                message=f"抱歉，{validation_errors[0] if validation_errors else '参数修改失败'}",
                intent_type="parameter_modification",
                parameter_changes={},
                explanation="\n".join(validation_errors) if validation_errors else "所有参数都超出有效范围",
                current_parameters={
                    "initial_state": current_initial_state,
                    "parameters": current_model_params
                },
                validation_errors=validation_errors,
                update_success=False,
                update_message="所有参数都超出有效范围",
                parameter_ranges=parameter_ranges
            )
    
    async def rainfall_hypothesis(self, question: str, session_id: str) -> Dict[str, Any]:
        """
        降雨假拟工具
        
        功能：
        - 解析用户的降雨假拟请求
        - 识别调整方法（按比例/平均分配）
        - 提取目标降雨量数值
        - 生成降雨假拟配置
        
        Args:
            question: 用户问题，如 "降雨按比例增加到50mm" 或 "降雨设为40mm、50mm和60mm"
            session_id: 会话ID
            
        Returns:
            符合 API 文档格式的字典
        """
        # LLM 解析用户意图，直接提取数值
        extraction_prompt = ChatPromptTemplate.from_messages([
            ("system", """你是降雨假拟参数提取专家。根据用户请求，提取降雨假拟的配置。

支持的调整方法：
1. "按比例调整" (proportional) - 按比例缩放降雨
2. "平均分配" (average) - 平均分配降雨量

用户可能说：
- "降雨按比例增加到50mm" → method: proportional, values: [50.0]
- "降雨平均分配到60mm" → method: average, values: [60.0]
- "降雨设为40mm、50mm和60mm" → method: proportional, values: [40.0, 50.0, 60.0]
- "降雨增加到39mm、58mm和70mm" → method: proportional, values: [39.0, 58.0, 70.0]

请直接提取用户说的数值，不要进行任何计算或转换。

请返回JSON格式，只返回JSON：
{{
  "method": "proportional",
  "values": [50.0]
}}

注意：
1. values 数组中的数值保留1位小数
2. 直接提取用户说的具体数值，不要计算百分比
3. method 只能是 "proportional" 或 "average"
4. 如果用户没有明确说明方法，默认使用 "proportional"
"""),
            ("human", "{question}")
        ])
        
        chain = extraction_prompt | self.llm
        response = await chain.ainvoke({"question": question})
        
        # 解析 LLM 返回的配置
        try:
            content = response.content.strip()
            # 去除可能的 markdown 标记
            if "```" in content:
                content = content.split("```")[1]
                if content.startswith("json"):
                    content = content[4:]
            
            rainfall_config = json.loads(content.strip())
            method = rainfall_config.get("method", "proportional")
            values = rainfall_config.get("values", [])
            
            # 确保 values 是浮点数且保留1位小数
            values = [round(float(v), 1) for v in values]
            
            if not values:
                return APIResponse.error_response(
                    message="无法提取降雨数值，请使用明确的数值，例如'降雨按比例增加到50mm'",
                    intent_type="rainfall_scenario",
                    errors=["无法提取降雨数值"]
                )
            
        except json.JSONDecodeError:
            return APIResponse.error_response(
                message="无法解析降雨配置",
                intent_type="rainfall_scenario",
                errors=["LLM返回格式错误"],
                explanation=f"LLM返回格式错误: {response.content[:100]}"
            )
        except Exception as e:
            return APIResponse.error_response(
                message=f"解析降雨配置失败: {str(e)}",
                intent_type="rainfall_scenario",
                errors=[str(e)],
                explanation=str(e)
            )
        
        # 格式化输出（按照 API 文档格式）
        method_name = "按比例调整" if method == "proportional" else "平均分配"
        
        # 生成消息
        if len(values) == 1:
            message = f"好的，我将采用【{method_name}】方式，目标降雨量为 [{values[0]}] mm"
        else:
            values_str = ", ".join([str(v) for v in values])
            message = f"好的，我将采用【{method_name}】方式，目标降雨量为 [{values_str}] mm"
        
        explanation = f"根据您的要求，使用{method_name}方式进行降雨假拟"
        
        return APIResponse.success_response(
            message=message,
            intent_type="rainfall_scenario",
            explanation=explanation,
            rainfall_simulation={
                "method": method,
                "values": values
            }
        )
    
    async def understand_and_execute(self, question: str, session_id: str) -> Dict[str, Any]:
        """
        问题理解器（类似 websocket_server.py 的 ask_with_llm）
        
        功能：
        1. LLM 分类器判断用户意图
        2. 调用对应工具（参数修改 或 降雨假拟）
        3. 返回结果
        
        Args:
            question: 用户问题
            session_id: 会话ID
            
        Returns:
            工具执行结果（符合 API 文档格式）
        """
        # 使用分类器 LLM（温度稍高一点）
        classifier_llm = ChatOpenAI(
            base_url=Config.VLLM_BASE_URL,
            model_name=Config.VLLM_MODEL_NAME,
            api_key=Config.VLLM_API_KEY,
            temperature=0.3,
            max_tokens=500,
        )
        
        # 分析用户问题，决定调用哪个工具
        classifier_prompt = ChatPromptTemplate.from_messages([
            ("system", """你是水文预报工具选择助手。根据用户问题，判断应该调用哪个工具：

工具1: modify_parameters_tool - 修改模型参数
  触发条件：用户想要修改参数，如 SA0, UA0, YA0, K, B, IMP 等
  典型问题："把SA0改成50"、"把K增加0.3"、"修改初始蓄水量"、"将SA0修改为10"

工具2: rainfall_hypothesis_tool - 降雨假拟
  触发条件：用户想要调整降雨量
  典型问题："降雨按比例增加到50mm"、"降雨设为40mm、50mm和60mm"、"降雨平均分配到60mm"

工具3: unknown - 未知意图
  触发条件：无法识别用户意图或与水文预报无关
  典型问题："今天天气怎么样"、"你好"、"帮我查一下资料"

请只回复工具名称（modify_parameters_tool / rainfall_hypothesis_tool / unknown），不要有其他内容。"""),
            ("human", "{question}")
        ])
        
        # 让 LLM 选择工具
        chain = classifier_prompt | classifier_llm
        response = await chain.ainvoke({"question": question})
        tool_name = response.content.strip().lower()
        
        # 调用相应工具
        if "unknown" in tool_name or "未知" in tool_name:
            # 未知意图
            return APIResponse.success_response(
                message="抱歉，我无法理解您的意图。您可以尝试以下方式：\n\n1. 参数修改：例如 \"将SA0修改为10\"\n2. 降雨假拟：\n   - 按比例调整：例如 \"降雨按比例增加到50mm\"\n   - 平均分配：例如 \"降雨平均分配到60mm\"",
                intent_type="unknown",
                understood=False
            )
        elif "rainfall" in tool_name or "降雨" in tool_name or "假拟" in tool_name:
            # 降雨假拟
            return await self.rainfall_hypothesis(question, session_id)
        else:
            # 参数修改
            return await self.modify_parameters(question, session_id)

